import requests
import aiohttp
import asyncio
import json


def search_papers(query: str, top_k: int =5):
    """
    根据文本查询搜索论文片段。
    """
    url = "http://180.184.65.98:38880/atomgit/search_papers"
    params = {'query': query, 'top_k': top_k}
    response = requests.get(url, params=params)

    if response.status_code == 200:
        try:
            json_data = response.json()
            return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"

async def asearch_papers(query: str, top_k: int = 5):
    """
    根据文本查询搜索论文片段的异步函数。
    """
    url = "http://180.184.65.98:38880/atomgit/search_papers"
    params = {'query': query, 'top_k': top_k}

    async with aiohttp.ClientSession() as session:
        async with session.get(url, params=params) as response:
            if response.status == 200:
                try:
                    json_data = await response.json()
                    return json_data
                except ValueError:
                    return "响应不是有效的 JSON 格式"
            else:
                return f"请求失败，状态码: {response.status}"

def query_by_paper_id(paper_id: str, top_k: int =5):
    """
    根据论文 ID 查询论文片段。需要精确匹配。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_paper_id"
    params = {
        'paper_id': paper_id,
        "top_k": top_k
        }
    response = requests.get(url, params=params)

    if response.status_code == 200:
        try:
            json_data = response.json()
            return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"


def query_by_title(title: str, top_k: int = 100):
    """
    根据论文标题查询论文片段。需要精确匹配。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_title"
    params = {
        'title': title,
        "top_k": top_k
        }
    response = requests.get(url, params=params)

    if response.status_code == 200:
        try:
            json_data = response.json()
            return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"

async def aquery_by_title(title: str, top_k: int = 100):
    """
    根据论文标题查询论文片段。需要精确匹配。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_title"
    params = {
        'title': title,
        "top_k": top_k
        }
    async with aiohttp.ClientSession() as session:
        async with session.get(url, params=params) as response:
            if response.status == 200:
                try:
                    json_data = await response.json()
                    return json_data
                except ValueError:
                    return "响应不是有效的 JSON 格式"
            else:
                return f"请求失败，状态码: {response.status}"

def metadata():
    """
    获取论文数据库的元数据信息。
    """
    url = "http://180.184.65.98:38880/atomgit/metadata"
    response = requests.get(url)

    if response.status_code == 200:
        try:
            json_data = response.json()
            # return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"


def query_by_title_contain(title: str, top_k: int = 1000):
    """
    搜索标题中包含特定文本的论文片段。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_title_contain"
    params = {
        'title': title,
        "top_k": top_k
        }
    response = requests.get(url, params=params)

    if response.status_code == 200:
        try:
            json_data = response.json()
            return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"


def query_by_chunk_contain(chunk: str, top_k: int = 1000):
    """
    搜索论文内容片段中包含特定文本的片段。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_chunk_contain"
    params = {
        'chunk': chunk,
        "top_k": top_k
        }
    response = requests.get(url, params=params)

    if response.status_code == 200:
        try:
            json_data = response.json()
            return json_data
            return json.dumps(json_data, indent=4)
        except ValueError:
            return "响应不是有效的 JSON 格式"
    else:
        return f"请求失败，状态码: {response.status_code}"

async def aquery_by_chunk_contain(chunk: str, top_k: int = 1000):
    """
    搜索论文内容片段中包含特定文本的片段。
    """
    url = "http://180.184.65.98:38880/atomgit/query_by_chunk_contain"
    params = {
        'chunk': chunk,
        "top_k": top_k
        }
    async with aiohttp.ClientSession() as session:
        async with session.get(url, params=params) as response:
            if response.status == 200:
                try:
                    json_data = await response.json()
                    return json_data
                except ValueError:
                    return "响应不是有效的 JSON 格式"
            else:
                return f"请求失败，状态码: {response.status}"

if __name__ == "__main__":

# print(type(search_papers("llm")) == list)
# print(type(query_by_paper_id(paper_id=1)) == list)
# print(type(query_by_title("Quantized Side Tuning: Fast and Memory-Efficient Tuning of Quantized Large Language Models")) == list)
# print(metadata())
# print(type(query_by_title_contain("llm")) == list)
# print(type(query_by_chunk_contain("llm")) == list)

# print(search_papers("In this paper, we highlight that both conformity and risk prefer", 1))
# print("===" * 30)
    print(query_by_paper_id(paper_id="6461b9c9d68f896efad43219"))
# print(query_by_title("Pessimism for Offline Linear Contextual Bandits Using $\\\\ell_p$ Confidence Sets."))
# print(query_by_title_contain("Which Matters Most"))
# print(query_by_chunk_contain("llm", 3))